package org.offsync.apps

import java.io.IOException
import java.net.InetAddress
import java.net.InetSocketAddress
import java.nio.ByteBuffer
import java.nio.channels.SelectionKey
import java.nio.channels.Selector
import java.nio.channels.SocketChannel
import java.nio.channels.spi.SelectorProvider
import java.util.Collections
import scala.util.control.Breaks

class NioClient(
    // The host:port combination to connect to
    private val hostAddress: InetAddress,
    private val port: Int) extends Runnable {
  // The selector we'll be monitoring
  private val selector = SelectorProvider.provider().openSelector()

  // The buffer into which we'll read data when it's available
  private val readBuffer = ByteBuffer.allocate(8192)

  // A list of PendingChange instances
  private val pendingChanges: java.util.List[ChangeRequest] = new java.util.LinkedList[ChangeRequest]()

  // Maps a SocketChannel to a list of ByteBuffer instances
  private val pendingData: java.util.Map[SocketChannel, java.util.List[ByteBuffer]] =
    new java.util.HashMap[SocketChannel, java.util.List[ByteBuffer]]()
  
  // Maps a SocketChannel to a RspHandler
  private val rspHandlers: java.util.Map[SocketChannel, RspHandler] =
    Collections.synchronizedMap(new java.util.HashMap[SocketChannel, RspHandler]())
  

  def send(data: Array[Byte], handler: RspHandler) = {
    // Start a new connection
    val socket: SocketChannel = this.initiateConnection()
    
    // Register the response handler
    this.rspHandlers.put(socket, handler)
    
    // And queue the data we want written
    this.pendingData.synchronized {
      var queue = this.pendingData.get(socket)
      if (queue == null) {
        queue = new java.util.ArrayList[ByteBuffer]()
        this.pendingData.put(socket, queue)
      }
      
      queue.add(ByteBuffer.wrap(data))
    }

    // Finally, wake up our selecting thread so it can make the required changes
    this.selector.wakeup()
  }

  def run() {
    while (true) {
      try {
        // Process any pending changes
        this.pendingChanges.synchronized {
          val changes = this.pendingChanges.iterator()
          
          val breaks = new Breaks
          breaks.breakable {
            while (changes.hasNext()) {
              val change = changes.next()
              change.`type` match {
              case ChangeRequest.CHANGEOPS =>
                val key = change.socket.keyFor(this.selector)
                key.interestOps(change.ops)
                breaks.break
              case ChangeRequest.REGISTER =>
                change.socket.register(this.selector, change.ops)
                breaks.break
              }
            }
          }
          
          this.pendingChanges.clear()
        }

        // Wait for an event one of the registered channels
        this.selector.select()

        // Iterate over the set of keys for which events are available
        val selectedKeys = this.selector.selectedKeys().iterator()
        while (selectedKeys.hasNext()) {
          val key = selectedKeys.next()
          selectedKeys.remove()

          if (key.isValid()) {
            // Check what event is available and deal with it
            if (key.isConnectable()) {
              this.finishConnection(key)
            } else if (key.isReadable()) {
              this.read(key)
            } else if (key.isWritable()) {
              this.write(key)
            }
          }
        }
      } catch { case e: Exception =>
        e.printStackTrace()
      }
    }
  }

  private def read(key: SelectionKey): Unit = {
    val socketChannel = key.channel().asInstanceOf[SocketChannel]

    // Clear out our read buffer so it's ready for new data
    this.readBuffer.clear()

    // Attempt to read off the channel
    var numRead: Int = 0
    try {
      numRead = socketChannel.read(this.readBuffer)
    } catch { case e: IOException =>
      // The remote forcibly closed the connection, cancel
      // the selection key and close the channel.
      key.cancel()
      socketChannel.close()
      return
    }

    if (numRead == -1) {
      // Remote entity shut the socket down cleanly. Do the
      // same from our end and cancel the channel.
      key.channel().close()
      key.cancel()
      return
    }

    // Handle the response
    this.handleResponse(socketChannel, this.readBuffer.array(), numRead)
  }

  private def handleResponse(
      socketChannel: SocketChannel,
      data: Array[Byte],
      numRead: Int) = {
    // Make a correctly sized copy of the data before handing it
    // to the client
    val rspData = new Array[Byte](numRead)
    System.arraycopy(data, 0, rspData, 0, numRead)
    
    // Look up the handler for this channel
    val handler = this.rspHandlers.get(socketChannel)
    
    // And pass the response to it
    if (handler.handleResponse(rspData)) {
      // The handler has seen enough, close the connection
      socketChannel.close()
      socketChannel.keyFor(this.selector).cancel()
    }
  }

  private def write(key: SelectionKey): Unit = {
    val socketChannel = key.channel().asInstanceOf[SocketChannel]

    this.pendingData.synchronized {
      val queue = this.pendingData.get(socketChannel)

      // Write until there's not more data ...
      val breaks = new Breaks
      breaks.breakable {
        while (!queue.isEmpty()) {
          val buf = queue.get(0)
          socketChannel.write(buf)
          if (buf.remaining() > 0) {
            // ... or the socket's buffer fills up
            breaks.break
          }
          
          queue.remove(0)
        }
      }

      if (queue.isEmpty()) {
        // We wrote away all data, so we're no longer interested
        // in writing on this socket. Switch back to waiting for
        // data.
        key.interestOps(SelectionKey.OP_READ)
      }
    }
  }

  private def finishConnection(key: SelectionKey): Unit = {
    val socketChannel = key.channel().asInstanceOf[SocketChannel]
  
    // Finish the connection. If the connection operation failed
    // this will raise an IOException.
    try {
      socketChannel.finishConnect()
    } catch { case e: IOException =>
      // Cancel the channel's registration with our selector
      System.out.println(e)
      key.cancel()
      return
    }
  
    // Register an interest in writing on this channel
    key.interestOps(SelectionKey.OP_WRITE)
  }

  private def initiateConnection(): SocketChannel = {
    // Create a non-blocking socket channel
    val socketChannel = SocketChannel.open()
    socketChannel.configureBlocking(false)
  
    // Kick off connection establishment
    socketChannel.connect(new InetSocketAddress(this.hostAddress, this.port))
  
    // Queue a channel registration since the caller is not the 
    // selecting thread. As part of the registration we'll register
    // an interest in connection events. These are raised when a channel
    // is ready to complete connection establishment.
    this.pendingChanges.synchronized {
      this.pendingChanges.add(new ChangeRequest(socketChannel, ChangeRequest.REGISTER, SelectionKey.OP_CONNECT))
    }
    
    return socketChannel
  }
}

object NioClient {
  def main(args: Array[String]): Unit = {
    try {
      val client = new NioClient(InetAddress.getByName("localhost"), 9876)
      val t = new Thread(client)
      t.setDaemon(true)
      t.start()
      val handler = new RspHandler()
      client.send("GET / HTTP/1.0\r\n\r\n".getBytes(), handler)
      handler.waitForResponse()
    } catch { case e: Exception =>
      e.printStackTrace()
    }
  }
}
